import numpy as np
import pymc as pm

from bokeh.plotting import figure, show
from bokeh.layouts import gridplot
from bokeh.models import ColumnDataSource, Span
from bokeh.io import output_notebook
output_notebook()

from tqdm import tqdm

R = 2 # num reps
G = 3 # num genes
K = 4 # num spike-in species

s_spike_true = np.array([[2, 20, 1000, 2000], [2, 20, 1000, 2000]])

with pm.Model() as model:
    # spike-in/transcriptome priors
    mu_b = pm.HalfNormal("mu_b", sigma = 1)
    #sigma_b = pm.LogNormal("sigma_b", mu = 0.1, sigma = 0.3)
    
    b_x = pm.HalfNormal('b_x', sigma = mu_b, shape = (R,G))
    b_s = pm.HalfNormal("b_s", sigma = mu_b, shape = (R,K)) 

    # idealized priors
    # hyperprior global alpha
    alpha_global = pm.Dirichlet('alpha_global', a = np.ones(G), shape = G)
    X_mrna = pm.LogNormal("X_mrna", mu = np.log(50000), sigma = 1, shape = R)
    conc = 500 # to give flexibility
    alpha = pm.Dirichlet('alpha', a = conc * alpha_global, shape = (R,G))

    # likelihood spike ins
    s_seq_obs = pm.Poisson("s_seq_obs", mu = s_spike_true * b_s)

    # likelihood transcriptome
    x_mrna = X_mrna[:, None] * alpha 
    x_seq_obs = pm.Poisson("x_seq_obs", mu = x_mrna * b_x)

    prior_pred = pm.sample_prior_predictive(samples = 1000)

Sampling: [X_mrna, alpha, alpha_global, b_s, b_x, mu_b, s_seq_obs, x_seq_obs]

with pm.Model() as model:
    # spike-in/transcriptome priors
    mu_b = pm.HalfNormal("mu_b", sigma = 1)
    sigma_b = pm.LogNormal("sigma_b", mu = 0.1, sigma = 0.1)
    
    b_x = pm.LogNormal('b_x', mu = mu_b, sigma = sigma_b, shape = (R,G))
    b_s = pm.LogNormal("b_s", mu = mu_b, sigma = sigma_b, shape = (R,K)) 

    # idealized priors
    # hyperprior global alpha
    alpha_global = pm.Dirichlet('alpha_global', a = np.ones(G), shape = G)
    X_mrna = pm.LogNormal("X_mrna", mu = np.log(2500), sigma = 1, shape = R)
    conc = 500 # to give flexibility
    alpha = pm.Dirichlet('alpha', a = conc * alpha_global, shape = (R,G))

    # likelihood spike ins
    s_seq_obs = pm.Poisson("s_seq_obs", mu = s_spike_true * b_s)

    # likelihood transcriptome
    #x_mrna = X_mrna[:, None] * alpha 
    x_mrna = pm.Multinomial("x_mrna", n = X_mrna, p = alpha)
    x_seq_obs = pm.Poisson("x_seq_obs", mu = x_mrna * b_x)

    prior_pred = pm.sample_prior_predictive(samples = 1000)

Sampling: [X_mrna, alpha, alpha_global, b_s, b_x, mu_b, s_seq_obs, sigma_b, x_mrna, x_seq_obs]

param = 'x_seq_obs' # x_seq_obs, s_seq_obs, alpha, b_x, b_s

prior_samples = prior_pred.prior[param].values
plots = []

# Loop through each replicate (2 reps)
for rep in range(R):
    for gene_ind in range(G):
        prior_gene = prior_samples[0, :, rep, gene_ind] 

        # set up histogram
        p1 = figure(title=f"Prior Predictive Check, {param}, Rep {rep+1}, Gene {gene_ind}",
            x_axis_label=param,y_axis_label="Count",
            width=400,height=300,background_fill_color="#fafafa")

        hist, edges = np.histogram(prior_gene, bins=50)

        p1.quad(
            top=hist,bottom=0, left=edges[:-1],right=edges[1:],
            fill_color="skyblue",line_color="black",legend_label="1000 random samples")

        plots.append(p1)

n_cols = 4
# Split into rows
grid = [plots[i:i+n_cols] for i in range(0, len(plots), n_cols)]
# Display
show(gridplot(grid))

prior_samples = prior_pred.prior['alpha_global'].values
plots = []

for gene_ind in range(G):
    prior_gene = prior_samples[0, :, gene_ind]  # get prior samples for this replicate and gene

    p1 = figure(title=f"Prior Predictive Check, alpha global, Gene {gene_ind}",
        x_axis_label="alpha-global",y_axis_label="Count",
        width=400,height=300,background_fill_color="#fafafa")
    
    hist, edges = np.histogram(prior_gene, bins=50)
    
    p1.quad(
        top=hist,bottom=0,left=edges[:-1],right=edges[1:],
        fill_color="skyblue",line_color="black", legend_label="1000 random samples")
 
    plots.append(p1)
n_cols = 4
# Split into rows
grid = [plots[i:i+n_cols] for i in range(0, len(plots), n_cols)]
# Display
show(gridplot(grid))

def simulate_data(mu_b_mu, X_mrna_mu, X_mrna_sigma, conc, G=3, K=4, R=2):
    """Simulate data for the hierarchical model."""
    mu_b = np.abs(np.random.normal(0, mu_b_mu))  
    b_x_true = np.abs(np.random.normal(0, mu_b, size=(R, G)))
    b_s_true = np.abs(np.random.normal(0, mu_b, size=(R, K)))

    alpha_global_true = np.random.dirichlet(alpha=np.ones(G))
    alpha_true = np.random.dirichlet(alpha=conc * alpha_global_true, size=R)
    
    X_mrna_true = np.random.lognormal(X_mrna_mu, X_mrna_sigma, size=R)
    x_mrna_true = np.array([
        np.random.multinomial(n=total, pvals=alpha_true[r])
        for r, total in enumerate(X_mrna_true)
    ])    #x_mrna_true = X_mrna_true[:, None] * alpha_true

    
    s_spike_true = np.array([[2, 20, 1000, 2000], [2, 20, 1000, 2000]])    
    x_seq_true = np.random.poisson(lam=b_x_true * x_mrna_true)
    s_seq_true = np.random.poisson(lam=b_s_true * s_spike_true)
    
    true_params = {
        "mu_b": mu_b,
        "b_x": b_x_true,
        "b_s": b_s_true,
        "alpha_global": alpha_global_true,
        "X_mrna": X_mrna_true,
        "alpha": alpha_true,
    }
    return s_spike_true, x_seq_true, s_seq_true, true_params

def run_model(s_spike_true, x_seq_true, s_seq_true, G, K, R, mu_b_mu, X_mrna_mu, X_mrna_sigma, conc):
    """Run the PyMC model and return trace and predictions."""
    with pm.Model() as model:
        # Spike-in/transcriptome priors
        mu_b = pm.HalfNormal("mu_b", sigma=mu_b_mu)
        b_x = pm.HalfNormal('b_x', sigma=mu_b, shape=(R, G))
        b_s = pm.HalfNormal("b_s", sigma=mu_b, shape=(R, K)) 
        
        # Idealized priors
        # Hyperprior global alpha
        alpha_global = pm.Dirichlet('alpha_global', a=np.ones(G), shape=G)
        alpha = pm.Dirichlet('alpha', a=conc * alpha_global, shape=(R, G))
        
        # Likelihood spike ins
        s_seq_obs = pm.Poisson("s_seq_obs", mu=s_spike_true * b_s, observed=s_seq_true)
        
        # Likelihood transcriptome
        X_mrna1 = pm.LogNormal("X_mrna1", mu=X_mrna_mu, sigma=X_mrna_sigma, shape=R)
        X_mrna = pm.Deterministic("X_mrna", pm.math.maximum(pm.math.floor(X_mrna1), 1))
       # x_mrna = X_mrna[:, None] * alpha 
        x_mrna = pm.Multinomial('x_mrna', n = X_mrna, p = alpha, shape = (R,G))
        x_seq_obs = pm.Poisson("x_seq_obs", mu=x_mrna * b_x, observed=x_seq_true)
        
        trace = pm.sample(1000, tune=1000, target_accept=0.90, progressbar = False)
        prior_pred = pm.sample_prior_predictive(samples=1000)
        ppc = pm.sample_posterior_predictive(trace)
    
    return trace, prior_pred, ppc

def compute_ranks(true_vals, posterior_samples):
    """Compute ranks for SBC analysis."""
    ranks = []
    N_sbc = true_vals.shape[0]
    
    for i in range(N_sbc):
        # Flatten param posterior samples for ith simulation
        posterior_i = posterior_samples[i]  # shape (n_samples, ...)
        # Flatten to 2D: (n_samples, param_dim)
        posterior_i_flat = posterior_i.reshape(posterior_i.shape[0], -1)
        true_i_flat = true_vals[i].reshape(-1)
        
        ranks_i = []
        for dim in range(len(true_i_flat)):
            rank = np.sum(posterior_i_flat[:, dim] < true_i_flat[dim])
            ranks_i.append(rank)
        ranks.append(ranks_i)
    
    return np.array(ranks)  # shape (N_sbc, param_dim)

def plot_individual_posteriors(true_vals, posterior_samples, param_name):
    """Plot individual posterior distributions for each SBC run, rep, and dimension."""
    N_sbc = len(posterior_samples)
    plots = []
    plot_titles = []
    
    for sbc_run in range(N_sbc):
        # Get shapes for this parameter
        post_shape = posterior_samples[sbc_run].shape  # (n_chains, n_samples, ...)
        true_shape = true_vals[sbc_run].shape if hasattr(true_vals[sbc_run], 'shape') else ()
        
        # Handle different parameter types
        if len(post_shape) == 2:  # Scalar parameter (n_chains, n_samples)
            post_samples = posterior_samples[sbc_run].flatten()
            true_val = true_vals[sbc_run]
            
            p = figure(title=f'{param_name} - Run {sbc_run+1}', width=250, height=200)
            hist, edges = np.histogram(post_samples, bins=30, density=True)
            # Ensure arrays have same length
            left_edges = edges[:-1]
            right_edges = edges[1:]
            p.quad(top=hist, bottom=0, left=left_edges, right=right_edges, 
                   fill_alpha=0.7, line_color="black", fill_color="skyblue")
            
            vline = Span(location=true_val, dimension='height', line_color='red', 
                        line_dash='dashed', line_width=2)
            p.add_layout(vline)
            p.xaxis.axis_label = 'Value'
            p.yaxis.axis_label = 'Density'
            plots.append(p)
            
        elif len(post_shape) == 3:  # 1D parameter (n_chains, n_samples, dim)
            for dim in range(post_shape[2]):
                post_samples = posterior_samples[sbc_run][:, :, dim].flatten()
                true_val = true_vals[sbc_run][dim]
                
                p = figure(title=f'{param_name}[{dim}] - Run {sbc_run+1}', width=250, height=200)
                hist, edges = np.histogram(post_samples, bins=30, density=True)
                # Ensure arrays have same length
                left_edges = edges[:-1]
                right_edges = edges[1:]
                p.quad(top=hist, bottom=0, left=left_edges, right=right_edges, 
                       fill_alpha=0.7, line_color="black", fill_color="skyblue")
                
                vline = Span(location=true_val, dimension='height', line_color='red', 
                            line_dash='dashed', line_width=2)
                p.add_layout(vline)
                p.xaxis.axis_label = 'Value'
                p.yaxis.axis_label = 'Density'
                plots.append(p)
                
        elif len(post_shape) == 4:  # 2D parameter (n_chains, n_samples, dim1, dim2)
            for rep in range(post_shape[2]):
                for gene in range(post_shape[3]):
                    post_samples = posterior_samples[sbc_run][:, :, rep, gene].flatten()
                    true_val = true_vals[sbc_run][rep, gene]
                    
                    p = figure(title=f'{param_name}[{rep},{gene}] - Run {sbc_run+1}', 
                              width=250, height=200)
                    hist, edges = np.histogram(post_samples, bins=30, density=True)
                    # Ensure arrays have same length
                    left_edges = edges[:-1]
                    right_edges = edges[1:]
                    p.quad(top=hist, bottom=0, left=left_edges, right=right_edges, 
                           fill_alpha=0.7, line_color="black", fill_color="skyblue")
                    
                    vline = Span(location=true_val, dimension='height', line_color='red', 
                                line_dash='dashed', line_width=2)
                    p.add_layout(vline)
                    p.xaxis.axis_label = 'Value'
                    p.yaxis.axis_label = 'Density'
                    plots.append(p)
    
    # Arrange plots in grid
    n_plots = len(plots)
    if n_plots == 0:
        print(f"No plots generated for {param_name}")
        return
        
    n_cols = min(4, n_plots)
    grid_plots = []
    for i in range(0, n_plots, n_cols):
        row = plots[i:i+n_cols]
        # Pad row with None if needed
        while len(row) < n_cols:
            row.append(None)
        grid_plots.append(row)
    
    grid = gridplot(grid_plots)
    show(grid)

def plot_mean_vs_true(true_vals, posterior_samples, param_name):
    """Plot posterior means vs true values."""
    mean_post = []
    true_flat = []
    
    for i in range(len(posterior_samples)):
        post_mean = posterior_samples[i].mean(axis=0)
        true_val = true_vals[i]
        
        if np.isscalar(post_mean):
            mean_post.append(post_mean)
            true_flat.append(true_val)
        else:
            mean_post.extend(post_mean.flatten())
            true_flat.extend(true_val.flatten() if hasattr(true_val, 'flatten') else [true_val])
    
    p = figure(title=f'Mean Posterior vs True for {param_name}', width=400, height=400)
    p.scatter(true_flat, mean_post, alpha=0.6, size=8)
    
    # Add diagonal line
    min_val = min(min(true_flat), min(mean_post))
    max_val = max(max(true_flat), max(mean_post))
    p.line([min_val, max_val], [min_val, max_val], color="red", line_dash="dashed", 
           line_width=2, legend_label='Perfect Recovery')
    
    p.xaxis.axis_label = 'True Value'
    p.yaxis.axis_label = 'Posterior Mean'
    p.legend.location = "top_left"
    
    show(p)

def plot_rank_histogram(ranks, param_name):
    """Plot rank histogram for SBC analysis."""
    ranks_flat = ranks.flatten()
    
    p = figure(title=f'Rank Histogram for {param_name}', width=600, height=300)
    
    max_rank = int(np.max(ranks_flat))
    hist, edges = np.histogram(ranks_flat, bins=max_rank + 1, range=(0, max_rank + 1))
    
    # Ensure arrays have same length
    left_edges = edges[:-1]
    right_edges = edges[1:]
    
    p.quad(top=hist, bottom=0, left=left_edges, right=right_edges, 
           fill_color="green", line_color="black", alpha=0.7)
    
    p.xaxis.axis_label = 'Rank'
    p.yaxis.axis_label = 'Frequency'
    
    show(p)

def plot_all_for_param(true_vals, posterior_samples, param_name):
    """Plot all SBC diagnostics for a parameter."""
    print(f"\n{'='*50}")
    print(f"SBC Results for parameter: {param_name}")
    print(f"{'='*50}")
    
    # Individual posterior distributions
    #print("\n1. Individual Posterior Distributions:")
    #plot_individual_posteriors(true_vals, posterior_samples, param_name)
    
    # Mean vs true scatter plot
    print("\n1. Posterior Mean vs True Value:")
    plot_mean_vs_true(true_vals, posterior_samples, param_name)
    
    # Rank histogram
    print("\n2. Rank Histogram:")
    ranks = compute_ranks(true_vals, posterior_samples)
    plot_rank_histogram(ranks, param_name)

def run_sbc(mu_b_mu, X_mrna_mu, X_mrna_sigma, conc, N_sbc=20, G=3, K=4, R=2):
    """Run Simulation-Based Calibration."""
    true_params_list = {
        'mu_b': [],
        'b_x': [],
        'b_s': [],
        'alpha_global': [],
        'X_mrna': [],
        'alpha': [],
    }
    posterior_samples = {
        'mu_b': [],
        'b_x': [],
        'b_s': [],
        'alpha_global': [],
        'X_mrna': [],
        'alpha': [],
    }
    
    for i in tqdm(range(N_sbc)):
        s_spike_true, x_seq_true, s_seq_true, true_param_dict = simulate_data(mu_b_mu, X_mrna_mu, X_mrna_sigma, conc, G, K, R)
        trace, prior_pred, ppc = run_model(s_spike_true, x_seq_true, s_seq_true, G, K, R, mu_b_mu, X_mrna_mu, X_mrna_sigma, conc)
                
        # Extract posterior samples for each parameter
        for param in posterior_samples.keys():
            posterior_samples[param].append(trace.posterior[param].values)
            true_params_list[param].append(true_param_dict[param])
            
    # Convert lists to arrays
    for param in posterior_samples.keys():
        posterior_samples[param] = np.array(posterior_samples[param])
        true_params_list[param] = np.array(true_params_list[param])
    
    return true_params_list, posterior_samples

# Suppress specific warning message
import warnings
warnings.filterwarnings("ignore", message='install "ipywidgets" for Jupyter support')

mu_b_mu = 3
X_mrna_mu = np.log(2500)
X_mrna_sigma = 1
conc = 500

N_sbc = 30
np.random.seed(42)

# Run SBC to get true params and posterior samples
print("Running SBC analysis...")
true_params, posterior_samples = run_sbc(mu_b_mu, X_mrna_mu, X_mrna_sigma, conc, N_sbc, G, K, R)

Running SBC analysis...

Multiprocess sampling (4 chains in 4 jobs)             | 0/30 [00:00<?, ?it/s]
CompoundStep
>NUTS: [mu_b, b_x, b_s, alpha_global, alpha, X_mrna1]
>Metropolis: [x_mrna]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 7 seconds.
There were 3566 divergences after tuning. Increase `target_accept` or reparameterize.
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
Sampling: [X_mrna1, alpha, alpha_global, b_s, b_x, mu_b, s_seq_obs, x_mrna, x_seq_obs]
Sampling: [s_seq_obs, x_seq_obs]

Multiprocess sampling (4 chains in 4 jobs)     | 1/30 [00:08<04:12,  8.70s/it]
CompoundStep
>NUTS: [mu_b, b_x, b_s, alpha_global, alpha, X_mrna1]
>Metropolis: [x_mrna]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 7 seconds.
There were 3577 divergences after tuning. Increase `target_accept` or reparameterize.
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
Sampling: [X_mrna1, alpha, alpha_global, b_s, b_x, mu_b, s_seq_obs, x_mrna, x_seq_obs]
Sampling: [s_seq_obs, x_seq_obs]

Multiprocess sampling (4 chains in 4 jobs)     | 2/30 [00:17<04:01,  8.62s/it]
CompoundStep
>NUTS: [mu_b, b_x, b_s, alpha_global, alpha, X_mrna1]
>Metropolis: [x_mrna]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 8 seconds.
There were 3564 divergences after tuning. Increase `target_accept` or reparameterize.
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
Sampling: [X_mrna1, alpha, alpha_global, b_s, b_x, mu_b, s_seq_obs, x_mrna, x_seq_obs]
Sampling: [s_seq_obs, x_seq_obs]

Multiprocess sampling (4 chains in 4 jobs)     | 3/30 [00:29<04:38, 10.32s/it]
CompoundStep
>NUTS: [mu_b, b_x, b_s, alpha_global, alpha, X_mrna1]
>Metropolis: [x_mrna]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 9 seconds.
There were 3579 divergences after tuning. Increase `target_accept` or reparameterize.
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
Sampling: [X_mrna1, alpha, alpha_global, b_s, b_x, mu_b, s_seq_obs, x_mrna, x_seq_obs]
Sampling: [s_seq_obs, x_seq_obs]

Multiprocess sampling (4 chains in 4 jobs)     | 4/30 [00:41<04:41, 10.81s/it]
CompoundStep
>NUTS: [mu_b, b_x, b_s, alpha_global, alpha, X_mrna1]
>Metropolis: [x_mrna]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 11 seconds.
There were 3566 divergences after tuning. Increase `target_accept` or reparameterize.
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
Sampling: [X_mrna1, alpha, alpha_global, b_s, b_x, mu_b, s_seq_obs, x_mrna, x_seq_obs]
Sampling: [s_seq_obs, x_seq_obs]

# import pickle

# # save results pickle
# filename = 'true_params, posterior_samples, run2'

# if not filename.endswith('.pkl'):
#     filename += '.pkl'
# sbc_data = {'true_params': true_params,
#            'posterior_samples':posterior_samples}
# with open(filename, 'wb') as f:
#     pickle.dump(sbc_data,f)

# For each parameter, plot SBC results
for param in true_params.keys():
    plot_all_for_param(true_params[param], posterior_samples[param], param)

==================================================
SBC Results for parameter: mu_b
==================================================

1. Posterior Mean vs True Value:

BokehUserWarning: ColumnDataSource's columns must be of the same length. Current lengths: ('x', 30), ('y', 30000)

2. Rank Histogram:

==================================================
SBC Results for parameter: b_x
==================================================

1. Posterior Mean vs True Value:

BokehUserWarning: ColumnDataSource's columns must be of the same length. Current lengths: ('x', 180), ('y', 180000)

2. Rank Histogram:

==================================================
SBC Results for parameter: b_s
==================================================

1. Posterior Mean vs True Value:

BokehUserWarning: ColumnDataSource's columns must be of the same length. Current lengths: ('x', 240), ('y', 240000)

2. Rank Histogram:

==================================================
SBC Results for parameter: alpha_global
==================================================

1. Posterior Mean vs True Value:

%load_ext watermark
%watermark -v -p numpy,bokeh,pymc,pandas,tqdm,warnings,jupyterlab

Run prior predictive checks and simulation-based calibration¶

Spike-in model with biological replicates¶

Imports¶

Model description¶

With biological replicates¶

Prior predictive checks¶

Simulation-based calibration¶

100 runs¶

Functions to simulate data from prior and run the model¶

Functions to compute rank and plot¶

Function to run entire SBC¶

Run entire SBC and plot¶